"""
Brief  : B样条操作，用于Robocon比赛路径规划
Author : 傅帅
Date   : 2020.4.14
"""
import math
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif']=['SimHei']  # 解决plt中文乱码

class BR_BSpline():
    """
    k             ： B样条阶次i，若要进行曲率半径计算，请保证k>=2
    controlPoints :  控制点[[x1,y1],[x2,y2],...]，其数目应大于等于k-1
    knots         :  节点，数目应严格等于‘控制点数目+阶次+1’，范围务必归一化在[0,1]，务必遵循Clamped规则
                     强烈建议不主动输入，使用程序自动生成的均布Clamped的节点
    example:
        k = 3
        cps = [[1,0],[1.5,1],[2,1.5],[3,2],[4,1],[3,0]]
        myBSp = BR_BSpline(k, cps)
        myBSp.draw()
    """

    def __init__(self, k, controlPoints, knots = []):
        """ 创建并初始化B样条操作实例 """
        self.k = k
        self.cps = np.array(controlPoints)
        if(len(self.cps) < k-1): raise ValueError("控制点数目太少，应大于等于k-1")
        auto_generate = True if len(knots) == 0 else False
        if auto_generate:  # 自动生成节点
            # 自动生成满足均布Clamped规则的节点
            knots_part1 = np.zeros(k)
            knots_part2 = np.linspace(0,1,len(self.cps)-k+1)
            knots_part3 = np.ones(k)
            self.knots = np.append(knots_part1, knots_part2)
            self.knots = np.append(self.knots, knots_part3)
        else:  # 输入节点
            if len(knots) != len(self.cps)+k+1: raise ValueError("knots数目应等于‘控制点数目+阶次+1’!")
            self.knots = knots
        self._basicFuncTbl = np.zeros([len(self.knots) - 1, self.k + 1])  # 基本函数表，节点个数-1行、阶数+1列
    
    def __call__(self, t):
        """
        根据曲线进程t计算曲线中某点坐标
        t: 曲线参数，取值范围务必为[0,1]。若为float型则返回对应坐标(x, y)，
           若为列表则返回坐标列表(x_list, y_list)
        """
        if type(t) == float or type(t) == np.float64:
            return self._calculate(t)
        elif (type(t) == list or type(t) == np.ndarray):
            x = list()
            y = list()
            for item in t:
                x.append(self._calculate(item)[0])
                y.append(self._calculate(item)[1])
            return(x,y)

    def draw(self, draw_radius = True):
        """
        画出该B样条的图像和曲率半径分布情况
        draw(draw_radius = False)可只画B样条曲线图，不画半径随t变化图
        """
        if len(self.cps) <= 2: raise ValueError("控制点少于三个时，请使用draw(draw_radius = False)")
        if draw_radius and self.k<2: raise ValueError("阶数k小于2时，请使用draw(draw_radius = False)")
        t_list = np.linspace(0, 1, 100)
        plt.figure(str(self.k)+"阶B样条", figsize = (18,9) if draw_radius else (9,9))
        # 曲线本体作图
        if draw_radius:
            plt.subplot("121")
        plt.title("B样条曲线图", fontsize="25")
        plt.xlabel("X坐标", fontsize="15")
        plt.xticks(fontsize="15")
        plt.ylabel("Y坐标", fontsize="15")
        plt.yticks(fontsize="15")
        x_list, y_list = self.__call__(t_list)
        plt.plot(x_list, y_list, 'b')
        cx_list = [p[0] for p in self.cps]
        cy_list = [p[1] for p in self.cps]
        plt.plot(cx_list, cy_list, 'or')
        plt.legend(["B-Spline", "ControlPoints"], fontsize="15")
        # 曲率半径作图
        if draw_radius:
            plt.subplot("122")
            plt.title("曲率半径随t变化图", fontsize="25")
            plt.xlabel("曲线参数t", fontsize="15")
            plt.xticks(fontsize="15")
            plt.ylabel("曲率半径R", fontsize="15")
            plt.yticks(fontsize="15")
            R_list = self.get_radius(t_list)
            plt.plot(t_list, R_list, 'm')
            plt.legend(["Radius"], fontsize="15")
        plt.show()

    def get_radius(self, t):
        """
        用法类似于__call__
        t: 曲线参数，取值范围务必为[0,1]。若为float型则返回对应半径R[m]，
           若为列表则返回半径列表R_list
        """
        if type(t) == float:
            return self._get_radius(t)
        elif (type(t) == list or type(t) == np.ndarray):
            R_list = list()
            for item in t:
                R_list.append(self._get_radius(item))
            return R_list

    def _get_radius(self, t):
        """
        计算某点对应半径
        t：float 
        return：radius
        """
        if self.k<2: raise ValueError("阶数k小于2时，无法计算曲率半径")
        if len(self.cps) <= 2: raise ValueError("控制点数小于等于2时，无法计算曲率半径")
        k = self.k
        cps = self.cps
        knots = self.knots
        # 计算该点的一阶导，方法参考某论文
        k1 = k-1  # 一阶导等价B样条的阶数
        knots1 = knots[1:-1]  # 一阶导等价B样条的节点
        cps1 = [k*(cps[i+1]-cps[i])/(knots[i+k+1]-knots[i+1]) for i in range(len(cps)-1)]  # 一阶导等价B样条的控制点
        Bsp1 = BR_BSpline(k1, cps1, knots1)  # 一阶导等价B样条
        # 计算该点的二阶导，方法参考某论文
        k2 = k1-1  # 二阶导等价B样条的阶数
        knots2 = knots1[1:-1]  # 二阶导等价B样条的节点
        cps2 = [k1*(cps1[i+1]-cps1[i])/(knots1[i+k1+1]-knots1[i+1]) for i in range(len(cps1)-1)]  # 二阶导等价B样条的控制点
        Bsp2 = BR_BSpline(k2, cps2, knots2)  # 二阶导等价B样条
        # 计算曲率半径和半径
        (x1, y1) = Bsp1(t)
        (x2, y2) = Bsp2(t)
        K = (x1*y2 - y1*x2) / math.sqrt((x1**2 + y1**2)**3)
        radius = 1/K
        return abs(radius)

    def _calculate(self, t):
        """
        根据曲线参数t计算坐标
        t：float 
        return：(x, y)
        """
        # 若t为0/1，后面公式会算出k阶基函数全为0，此处直接处理此特殊情况
        if t == 0: return self.cps[0]
        if t == 1: return self.cps[-1]
        # 正常情况↓
        knots = self.knots
        tmpTbl = self._basicFuncTbl.T  # 为方便遍历，先转置，tmpTbl中的行为基函數表的列
        # 根据t更新基函数表
        for deg in range(len(tmpTbl)):
            if deg == 0:  # 基函数表第一列，计算0阶初始值
                for i in range(len(knots) - 1):  # 此处i表示节点下标
                    tmpTbl[deg][i] = 1 if knots[i] <= t <= knots[i+1] else 0
            else:  # 其他情况递归计算
                for i in range(len(knots) - 1):
                    try:
                        # 此处规定0/0 = 0，foo表示基函数表计算公式中的各块
                        foo1 = (t - knots[i]) * tmpTbl[deg-1][i]
                        foo2 = (knots[i+deg] - knots[i])
                        if foo1 == 0 and foo2 == 0: part1 = 0
                        else: part1 = foo1/foo2

                        foo3 = (knots[i+deg+1] - t) * tmpTbl[deg-1][i+1]
                        foo4 = (knots[i+deg+1] - knots[i+1])
                        if foo3 == 0 and foo4 == 0: part2 = 0
                        else: part2 = foo3/foo4

                        tmpTbl[deg][i] = part1 + part2
                    except IndexError:  # 和最终计算无关的部分基函数表项，在try中会有knots[i+deg+1]越界，直接置为-1
                        tmpTbl[deg][i] = -1
        self._basicFuncTbl = tmpTbl.T  # 转置回去

        # 根据基函数表计算分配权重，计算最终结果
        p = [0.,0.]
        weight_sum = self._basicFuncTbl[:len(self.cps)].sum(0)[-1]  # 所有k阶基函数的和，用于权重归一化
        for i in range(len(self.cps)):
            p += self.cps[i] * self._basicFuncTbl[i][-1] / weight_sum  # 基函数表每行最后一列为该控制点的权重
        return p

    def path_generate_dense(self, max_a, max_vel, slice_num = 1000):
        """
        密集路径生成，计算并返回[[t, x, y, R, a_t(能提供的最大切向a), vel, s, time, vel_lim], ...]的列表，单位均为mm
        max_a       : 最大加速度(m/s^2)
        max_vel     : 最大速度(m/s)
        slice_num   : 将[0,1]切片数，应取较大值
        -------------------------------------------------------------------------------------------------------
        @attention  : 暂不支持两控制点的直线，也不想加了，需要直线的话自己用三个控制点拖吧
        """
        HYPERPARAM = 0.95  # 表示反向递推时的起始速度 = HYPERPARAM * vel_lim，以保证留出一定切向加速度用于加减速
        max_a *= 1000.0
        max_vel *= 1000.0
        path_tbl = np.zeros([slice_num, 9])
        t_tbl = np.linspace(0, 1, slice_num)
        # 初始化一些值
        for i in range(slice_num):
            path_tbl[i][0] = t_tbl[i]
            path_tbl[i][1], path_tbl[i][2] = self._calculate(t_tbl[i])
            path_tbl[i][3] = self._get_radius(t_tbl[i])  # 曲率半径
            # 在该点可通过的最大速度，为避免浮点型精度损失导致在算切向速度时max_a^2-[由vel_lim和R算出a_n^2]小于零的问题，乘0.999
            path_tbl[i][8] = math.sqrt(max_a * path_tbl[i][3])*0.999  
        # 整体处理思路：未解决不确定何时开始加减速的问题，以速度限制函数图像的波谷为分界，将整条路径分解成若干小段。
        #           然后从头开始对每一个小段从两端分别进行前向和反向递推，处理合成为一个完整的小段，再将各小段拼接。
        #           第一遍遍历先计算速度和可提供的最大切向加速度，此二者确定之后最后再整理路程和时间
        path_tbl_split = []  # 存储小段
        last_valley_index = 0  # 上个谷底下标
        for i in range(slice_num-2):
            index = i+1
            if path_tbl[index][8] < path_tbl[index-1][8] and path_tbl[index][8] <path_tbl[index+1][8]:  # 若在波谷
                path_tbl_split.append(path_tbl[last_valley_index:index])
                last_valley_index = index
        path_tbl_split.append(path_tbl[last_valley_index:len(path_tbl)])  # 加入最后一段
        # 遍历所有小段block
        for index, path_block in enumerate(path_tbl_split):
            # 前向递推，先确定初值
            if path_block[0][0] == 0:  # 若为第一个小段
                path_block[0][4] = max_a
                path_block[0][5] = 0
            else:  # 正常情况下由前一点递推得到
                p_prev = path_tbl_split[index-1][-1]
                p_initial = path_block[0]
                delta_s = math.sqrt((p_initial[1] - p_prev[1])**2 + (p_initial[2] - p_prev[2])**2)
                V1 = math.sqrt(2*p_prev[4]*delta_s + p_prev[5]**2)
                p_initial[5] = V1 if V1 < p_initial[8] else p_initial[8]
                p_initial[5] = p_initial[5] if p_initial[5] < max_vel else max_vel
                a_n = (p_initial[5]**2) / p_initial[3]
                p_initial[4] = math.sqrt(max_a**2 - a_n**2)
            # 前向递推，先直接操作path_block，若需进行反向递推，则再覆盖即可
            for i in range(len(path_block)-1):  # 假设以上个点的加速度进行匀加速，不断计算下个点
                p1 = path_block[i]
                p2 = path_block[i+1]
                delta_s = math.sqrt((p2[1] - p1[1])**2 + (p2[2] - p1[2])**2)  # 位移
                V1 = math.sqrt(2*p1[4]*delta_s + p1[5]**2)  # 两点间极限加速下的V1 = sqrt(2ax + V0^2)
                if V1 > p2[8]: 
                    # 若会超过速度限制，前向递推可到此结束，原理可画图理解
                    break  
                p2[5] = V1 if V1 < p2[8] else p2[8]  # 若会超过向心加速度造成的速度限制，vel改为vel_lim
                p2[5] = p2[5] if p2[5] < max_vel else max_vel  # 若超过路径要求的最大速度，设为最大速度
                a_n = (p2[5]**2) / p2[3]  # 向心加速度 = V^2/R
                p2[4] = math.sqrt(max_a**2 - a_n**2)  # max_a提供切向和法向加速度
            # 若前向递推一直以极限性能加速，却仍达不到速度限制，则认为结果可用(最后一段除外，必须进行反向递推)
            if path_block[-1][5] <= path_block[-1][8] and path_block[-1][5] != 0 and index != len(path_tbl_split)-1:
                pass
            else:  # 若最后一点为0或者超过vel_lim，即前向递推未能一直算完，则需要进行反向递推
                # 反向递推，目的为不断反向更新path_block，直至找到应开始全力减速的点 
                # 若当前path_block更新至开头仍未得到最终合理曲线，则对上一个path_block重复操作
                partial_finish_flag = 0  # 表示是否得到合理曲线
                _path_block = path_block  # 表示当前在进行反向递推的block
                num = 1  # 表示这一轮反向递推重复操作的次数，一般为1
                while partial_finish_flag == 0:  # 循环直到得到合理曲线
                    # 确定初值
                    if num == 1:  # 第一次反向递推
                        if index != len(path_tbl_split)-1:  # 正常情况下根据vel_lim确定起始速度
                            V = _path_block[-1][8] * HYPERPARAM
                            _path_block[-1][5] = V if V < max_vel else max_vel
                            _path_block[-1][4] = _path_block[-1][5]**2 / _path_block[-1][3]
                        else:  # 最后一段起始速度为0
                            _path_block[-1][5] = 0
                            _path_block[-1][4] = max_a
                    else:  # 上次反向递推推完整个block都没得到合理曲线，则依照上次末尾值递推
                        p_prev = path_tbl_split[index-num+2][0]  # 上次反向递推的block的第一个点
                        p_initial = _path_block[-1]
                        delta_s = math.sqrt((p_initial[1] - p_prev[1])**2 + (p_initial[2] - p_prev[2])**2)
                        V1 = math.sqrt(2*p_prev[4]*delta_s + p_prev[5]**2)
                        V1 = V1 if V1 < p_initial[8] else p_initial[8]
                        p_initial[5] = V1 if V1 < max_vel else max_vel
                        a_n = (p_initial[5]**2) / p_initial[3]
                        p_initial[4] = math.sqrt(max_a**2 - a_n**2)
                    # 反向递推
                    for i in range(len(_path_block)-1):
                        p1 = _path_block[-i-1]
                        p2 = _path_block[-i-2]
                        delta_s = math.sqrt((p2[1] - p1[1])**2 + (p2[2] - p1[2])**2)  # 位移
                        V1 = math.sqrt(2*p1[4]*delta_s + p1[5]**2)  # 两点间极限加速下的V1 = sqrt(2ax + V0^2)
                        if p2[5] != 0 and V1 > p2[5]:  # 当前点速度不为零而且已经低于反向递推的速度
                            # 表示此区块内反向递推与前向递推的速度曲线已经相交，得到了合理的最终速度曲线
                            # 此点以及之前为全力加速部分，此点以及之后为全力减速部分
                            partial_finish_flag = 1
                            break
                        V1 = V1 if V1 < p2[8] else p2[8]
                        p2[5] = V1 if V1 < max_vel else max_vel  # 若超过路径要求的最大速度，设为最大速度
                        a_n = (p2[5]**2) / p2[3]  # 向心加速度 = V^2/R
                        p2[4] = math.sqrt(max_a**2 - a_n**2)  # max_a提供切向和法向加速度
                    _path_block = path_tbl_split[index-num]  # 若此次反向递推未能得到最终合理曲线，则对上一个block重复操作
                    num += 1
        # 此时，整条路径已经计算完成，接下来进行整理，以及s和t的计算
        tmp = []
        for block in path_tbl_split:
            for point in block:
                tmp.append(point)
        path_tbl = np.array(tmp)  # 重新将path_tbl_split中的blokc整理进path_tbl
        # 计算s和t
        path_tbl[0][6] = 0  # 初始s为0
        path_tbl[0][7] = 0  # 初始t为0
        for i in range(len(path_tbl)-1):
            p1 = path_tbl[i]
            p2 = path_tbl[i+1]
            delta_s = math.sqrt((p2[1] - p1[1])**2 + (p2[2] - p1[2])**2)  # 位移
            p2[6] = p1[6] + delta_s
            delta_t = delta_s/(p1[5]/2 + p2[5]/2)
            p2[7] = p1[7] + delta_t
        # 调试用画图，画出最终的vel_lim - t 和 vel - t
        # plt.plot(path_tbl[:,0],path_tbl[:,-1])
        # plt.plot(path_tbl[:,0],path_tbl[:,5])
        # plt.legend(["vel_lim", "vel"], fontsize="15")
        # plt.show()
        return path_tbl

if __name__ == "__main__":
    # test
    k = 3
    cps = [[1,0],[1.5,1],[2,1.5],[3,2],[4,1],[3,0]]
    myBSp = BR_BSpline(k, cps)
    myBSp.draw()